# train.py

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    # model = YOLO('yolov8n.yaml')  # 从YAML创建一个新模型

    # 加载预训练模型(推荐用于训练)
    # model = YOLO('D:/pythonProject/weed_detaction/ultralytics-main/yolov8n.pt')
    # 不加载加载预训练模型
    model = YOLO('./ultralytics/cfg/models/v8/yolov8.yaml')
    # 训练改进后的网络，并使用预训练模型
    # model = YOLO('D:/pythonProject/weed_detaction/ultralytics-main/ultralytics/cfg/models/v8/yolov8-DSConv.yaml').load('D:/pythonProject/weed_detaction/ultralytics-main/yolov8n.pt')
    # 训练改进后的网络，不使用预训练模型
    # model = YOLO('D:/pythonProject/weed_detaction/ultralytics-main/ultralytics/cfg/models/v8/yolov8-DSConv.yaml')
    # 将sppf换成注意力机制
    # model = YOLO('D:/pythonProject/weed_detaction/ultralytics-main/ultralytics/cfg/models/v8/yolov8-Linear-ATT.yaml')


    # 另一电脑运行
    # model = YOLO('D:/jkd/pythonProject/ultralytics-main/ultralytics/cfg/models/v8/yolov8n.yaml')
    # 训练模型
    model.train(data="my_data.yaml",
                imgsz=640,  # 输入图像的大小为整数或 w,h
                epochs=150,  # 要训练的次数
                batch=8,  # 每批次的图像数量（AutoBatch 为 -1）
                  # 要运行的设备，即 cuda device=0 或 device=0,1,2,3 或 device=cpu
                workers=4,  # 用于数据加载的工作线程数（如果是 DDP，则为每个 RANK）
                lr0=0.01,
                resume=True)  # True的时候则从上一个检查点恢复训练